Multi-engine backends: pushdown DuckDB adapter, engine-neutral Arrow dataset, caching helpers - #227
Open
Mmoncadaisla wants to merge 99 commits into
Open
Multi-engine backends: pushdown DuckDB adapter, engine-neutral Arrow dataset, caching helpers#227Mmoncadaisla wants to merge 99 commits into
Mmoncadaisla wants to merge 99 commits into
Conversation
Prototype of the "xarray-sql as the Xarray <-> engine translator" idea: the library owns two seams and nothing else. Seam 1 (register) attaches a lazy Dataset as a table on an engine's own connection; seam 2 (round-trip) turns any engine's Arrow result plus a template Dataset back into a labeled xr.Dataset. Dialects, geometry, H3, and optimizers stay with each engine and its extension ecosystem. - xarray_sql/backends/: adapter protocol + dispatch on connection type; DataFusion adapter delegates to the existing table provider, DuckDB adapter registers a re-scannable Arrow C-stream view (fresh lazy reader per scan: lazy AND re-queryable, no pushdown yet). - xarray_sql/roundtrip.py: engine-agnostic to_dataset() accepting DuckDB relations, pyarrow Tables/readers, or any __arrow_c_stream__ object; reuses the batches->Dataset core extracted from ds._materialize. - xql.register / xql.to_dataset exported at top level; [duckdb] extra. - docs/engines.md: the engine model, DuckDB usage, relation to duckdb-zarr (engine-native Zarr path; this adapter covers the rest of the xarray reader surface plus the labeled round-trip). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the default DuckDB registration object with XarrayPushdownDataset, a pyarrow.dataset.Dataset subclass (the pattern Lance uses for LanceDataset). DuckDB classifies it by isinstance and calls scanner(columns=..., filter=...) once per query, enabling: - projection pushdown: only the data variables a query mentions are loaded from storage; - chunk pruning: per-dimension shadow FileSystemDataset fragments carry each chunk's coordinate range as a partition_expression, so Arrow's guarantee simplification decides satisfiability for any predicate shape with no expression parsing on our side. One shadow per dimension keeps fragment counts at sum(n_d) instead of prod(n_d); - parallel production: surviving chunks are loaded by a bounded prefetch thread pool. DuckDB deletes pushed comparison conjuncts from its plan and never re-applies them, so the scanner always applies the exact expression via pyarrow Scanner; pruning is only an optimization. Filter-only columns absent from the projection are discovered by probing the expression against an empty table and widening on the miss. 10M-row benchmark vs the v1 stream: full scan 0.52s -> 0.035s, 1% time-filtered scan 0.24s -> 0.006s. A native-resolution bounding-box GROUP BY over a 9.13B-pixel cloud GeoTIFF answers in 0.8s on a plain duckdb connection (previously minutes). XarrayArrowStream stays as the dependency-light no-pushdown fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Bucketed two-level shadow pruning: axes with more chunks than the 1024-fragment fanout get a coarse bucket shadow refined lazily per surviving bucket, bounding pruning cost for finely partitioned datasets (e.g. hourly-chunked reanalysis time axes with hundreds of thousands of chunks). Refinement is skipped when a predicate keeps most buckets, where it cannot pay for itself. - Shadow datasets carry the full table schema so compound predicates referencing several columns bind and prune correctly. - Mixed-dimension datasets split into one table per dimension group (<name>_<dims>), sharing a single read of the dimension coordinates across sub-tables. - register() forwards adapter-specific kwargs (batch_size, prefetch on DuckDB; table_names on DataFusion). - Edge cases covered by tests: fully-pruned empty scans, LIMIT early termination, descending coordinate axes, uint8/string variables, 5000-chunk bucketed pruning loading exactly one chunk, kwarg forwarding. - duckdb extra pinned to >=1.4 (semantics verified against 1.5.4); engines doc gains production notes; pushdown benchmark added to benchmarks/. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Coordinate columns of a C-ordered partition are exact repeat/tile patterns of the dimension coords: dim k's flat column is its values each repeated prod(shape[k+1:]) times, tiled prod(shape[:k]) times. Building each column once per partition with those two sequential-write kernels and emitting batches as zero-copy Arrow slices replaces the per-batch division/modulo plus gather, making the pivot ~2.9x faster (53M -> 154M rows/s on a 10M-row, 3-dim dataset) with bitwise-identical output. Both engines benefit: iter_record_batches feeds the DataFusion table provider and the DuckDB pushdown scanner alike. The fast path holds the partition's full coordinate columns in memory (rows x 8 bytes x n_dims), so it is gated at 8M rows; larger partitions (e.g. single-time-step reanalysis) keep the O(batch_size) streaming path. The memory-profile test bands move accordingly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
rioxarray serializes GDAL reads behind a lock by default, capping any scan at single-stream speed regardless of the adapter's prefetch pool. lock=False measured 6x on full scans of a 9-billion-pixel cloud GeoTIFF (277s -> 43s) and makes remote reads as fast as a local copy. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Move XarrayPushdownDataset and XarrayArrowStream to
backends/pyarrow.py: the pushdown dataset is a real
pyarrow.dataset.Dataset, so it serves every consumer of that protocol,
not just DuckDB. The new public constructor xql.arrow_dataset(ds)
makes that explicit:
pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) # Polars, lazy
con.register("t", xql.arrow_dataset(ds)) # DuckDB
xql.arrow_dataset(ds).to_table(columns=..., filter=...) # pyarrow
Verified with Polars 1.42: scans are lazy, predicates and projections
push into the dataset (a filtered group-by read 1 of 20 chunks and 3
of 5 columns), results match xarray exactly, and Polars frames
round-trip through xql.to_dataset via the Arrow PyCapsule protocol.
The DuckDB adapter is now a thin registration shim over the shared
module. polars added to the test extra.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
get_fragments(filter=...) yields one fragment per source chunk, pruned by the same per-dimension shadow index the scanner uses. This is the protocol datafusion-python's register_dataset consumes (one DataFusion partition per fragment, filter pushdown marked Exact — safe because fragment scanners apply the expression row-exactly) and enables the Dask pattern from_map(lambda f: f.to_table().to_pandas(), ds.get_fragments()). Fragments carry a __dask_tokenize__ hook since the parent dataset is deliberately unpicklable. One arrow_dataset() object now serves DuckDB, Polars, DataFusion, Dask, and pyarrow itself. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two fast paths in the batches->Dataset core (shared by every engine's round-trip): - Uniformly spaced axes (rasters, regular time steps, ascending or descending) locate each row with rint((value - origin) / step) instead of a per-row binary search plus argsort remap. 2.2x on a 25M-row chunk-ordered window (0.96s -> 0.43s); axes that are not affine within a quarter step keep the searchsorted path. - Results that form the complete grid in C order (ORDER BY'd scans, single-chunk windows) skip positional scatter entirely and reshape the value column. Both are detected, never assumed; sparse and arbitrarily ordered results fall back to the general scatter. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Registered xarray tables are virtual: every query re-streams the source. For statistics asked repeatedly, xql.materialize(con, name, query, order_by=...) pays the scan once into a native engine table (sorted so coordinate columns compress and zone maps prune), and xql.pyramid(con, name, table, aggs=..., base_cell=..., levels=...) builds a CARTO-tileset-style multi-resolution pre-aggregated cube: level 0 bins the source in a single pass, coarser levels roll up from the level below, so any zoom/extent query is a range scan over a small table. Aggregates are restricted to decomposable kinds (sum/count/min/max) so roll-ups stay exact; averages derive from sum + count at query time. Both helpers dispatch through a new EngineAdapter.run_sql seam and are tested identically on DuckDB connections and DataFusion contexts (the SQL they emit is restricted to the shared dialect subset; DataFusion INSERT is lazy, so run_sql collects). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Measured guidance in one place: parallel source reads (LIBERTIFF / lock=False, 11x on full raster scans; zarr async.concurrency), chunk sizing for the scan, adapter knobs (prefetch, batch_size), what pushdown does and does not cover, materialize/pyramid for repeated statistics, and the round-trip fast paths. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- NaN/NaT in a chunk's dimension coordinate poisoned the shadow guarantee into (dim >= NaN), which Arrow simplifies every predicate against as false — chunks with matching rows were silently pruned. Such spans now carry an always-true guarantee (unprunable). - _affine_axis accepted NaN-poisoned axes (NaN > tol is false), letting the affine scatter place values in wrong cells when a result's dim column contains NULLs. The acceptance test is now a <= .all() so NaN rejects and the positional searchsorted path handles it. - Projected scans omit dimension columns from the scan schema; the cftime coordinate-conversion loop assumed every dim had a schema field and raised KeyError on e.g. SELECT SUM(v) over a 360_day dataset. Dims absent from the schema are now skipped. - pyramid() rebinned float cell origins level-over-level, occasionally aliasing boundary points into a neighboring cell. Cells are now tracked as integer indices that halve exactly per level, with float origins kept as query labels. Also: docs notes from stress testing (Polars is_in float-literal pushdown caveat, per-thread DuckDB cursor re-registration pattern). Regression tests for all four fixes; battle-test matrix passed on duckdb 1.4.5 LTS and 1.5.4, 4-thread concurrency, dtype zoo, and a 600-query memory soak. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mmoncadaisla
marked this pull request as draft
July 14, 2026 06:03
Polars (>= 1.40) passes batch_size through Dataset.to_batches to size its streaming-engine morsels; the scanner accepted the kwarg and dropped it, so consumers had no control over batch granularity. The emitted batch size matters beyond Polars: engines parallelize per record batch (DuckDB slices each batch into 2048-row vectors across threads, Acero schedules one filter task per batch), so it is the downstream-parallelism granule. Also pins two consumer-contract properties with tests: batch_size reaches the emitted batches through both scanner() and to_batches(), and the table schema never contains view types — one view-typed column disables DuckDB filter pushdown for the entire table (duckdb-python#227). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
count_rows() previously routed through a full scanner() scan, pivoting every surviving chunk to count rows the chunk grid already knows. Three-way split instead: pruned chunks contribute nothing, chunks whose coordinate ranges PROVE the filter true contribute their size arithmetically (no I/O), and only undecided boundary chunks are scanned — reading just the columns the filter references. Strictness is decided by Arrow itself, no expression parsing: a chunk with conjunctive guarantee G satisfies filter F everywhere iff G AND NOT F is unsatisfiable, which get_fragments(filter=~F) over per-chunk guarantee fragments answers (the Iceberg inclusive/strict evaluator pattern). Everything undecidable — NaN spans, data-variable predicates, >4096 survivors, unsupported expression shapes — falls back conservatively to an exact boundary scan. Also fixes columns=[] being treated as "all columns" (falsy-list bug): an explicitly empty projection is now a real projection, served by zero-column batches whose row counts are chunk arithmetic. Measured (10M-row synthetic, time chunked by 10): - unfiltered count: full scan -> 0.01 ms, zero chunks read - 5-day mid-chunk time range: exact count in 19 ms reading only the 2 boundary chunks (11 interior chunks proven arithmetically) - data-variable filter (no guarantee): still row-exact, all chunks scanned as before Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Finely chunked sources (an hourly-stepped time axis is one chunk per hour) pay one store round-trip per surviving chunk. With coalesce_rows=N, runs of consecutive surviving chunks along the most finely chunked dimension are merged into single isel reads of at most N rows, after pruning and per-run (no gaps are ever read). On Zarr sources the merged read fetches its member chunks through the store's own concurrent batching instead of one request per chunk through the prefetch pool. Scanner path only: get_fragments() keeps one fragment per source chunk so fragment consumers (DataFusion, dask) retain their parallelism granularity. Measured on ARCO-ERA5 over anonymous GCS (1.32M hourly time chunks, prefetch=16, coalesce_rows=8M), identical results both ways: - 1 day x Iberia bbox: 2.23s -> 1.18s (24 reads -> 4) - 1 week x full globe (174M rows): 9.82s -> 6.62s (168 reads -> 24) Peak RSS grows with prefetch x merged-block size as documented (0.7 GB -> 1.2 GB here); size coalesce_rows/prefetch together. Off by default: memory scales with the merged block, and local or coarsely chunked sources gain nothing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Each scan previously created (and tore down) its own ThreadPoolExecutor. Beyond the per-scan setup cost, spawning OS threads from inside an engine's scan callback is exactly the embedded-engine hazard the ecosystem keeps rediscovering (pg_duckdb serializes all host calls; ParadeDB routes engine work through dedicated pools): thread startup contends with concurrent.futures' process-global shutdown lock and with the consumer's own pool management, observed as intermittent deadlocks when scans are driven from dask worker threads. The pool now lives on the dataset, its threads started at construction time — never inside an engine callback. Scans that stop early (LIMIT) cancel their queued loads instead of shutting the pool down. Single-block scans (a lazy round-trip window that maps onto one source chunk) skip the pool entirely. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
xql.to_dataset gains chunks= and coords=: data variables are reconstructed window by window on access, each window re-executing the engine's query narrowed to its coordinate range with the engine's own typed expression API (never rendered SQL text). Over a table registered through xarray-sql, the window's range predicate flows back into chunk pruning at the source, so accessing one output chunk reads only the source chunks it maps onto. The engine-specific surface of the existing DataFusion lazy path (expression building, per-dim distinct discovery, schema access) is extracted into LazyResultHandle implementations in the new xarray_sql.lazyscan module; SQLBackendArray and _build_lazy_scan are now engine-neutral, and the DataFusion wrapper path is byte-identical in behavior (39 pre-existing round-trip tests unchanged and green). Contiguous window requests become two-literal range predicates (engines can prune on them); stepped/fancy indexers fall back to explicit value lists, exact by construction. Engine support: - Polars LazyFrame/DataFrame: full chunked support; per-window fetches run on the streaming engine. Verified deadlock-free under threaded dask (6/6 stress runs) and correct on descending coordinates, stepped indexers, filtered and aggregated queries. - DataFusion DataFrame: unchanged, now also reachable through the engine-agnostic xql.to_dataset. - DuckDB relations: eager round-trip fully supported through a dedicated single engine thread (concurrent materialization of derived relations corrupts shared pending-query state — verified — so all handle calls are funnelled through one thread). Chunked reconstruction FAILS FAST with guidance instead of hanging: re-executing a relation that scans a Python-backed table while other threads start/stop deadlocks intermittently inside duckdb-python/CPython (~50% of runs, macOS/CPython 3.12; persists with SET threads=1, connection serialization, and pool pre-warming; Polars is clean under the identical topology). Reproducer kept in the working notes for an upstream report. coords="template" skips per-dim DISTINCT discovery when the result spans the template's full extent, making construction free of source reads for unfiltered scans on any engine. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Documents which engines support chunked reconstruction and why DuckDB relations fail fast (upstream duckdb-python deadlock, reproduced and isolated); states the memory bound (prefetch x pivoted-block-size) with the ARCO-ERA5 measurements that back it, and the coalesce_rows memory/latency tradeoff. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Asserts the SHAPE of the work — exactly which source chunks each query reads (via the scanner's iteration callback) and exact row counts — not just answers or timings, so a pruning/coalescing/fast-path regression that silently falls back to scanning everything fails loudly (the plan-shape-assertion pattern; benchmark-suite tripwires caught regressions plain tests missed in comparable projects). Measured this run (anonymous GCS, 1.32M-chunk hourly time axis): - day+bbox: 2.8s / exactly 24 chunk reads (4 reads coalesced, 1.6s) - week globe (174M rows): 16.6s / exactly 168 reads - count(*) over January (772M rows): 0.09s / ZERO reads (arithmetic) - Polars lazy round-trip: construction 0 reads; 1-day window compute reads exactly its 4 coalesced blocks Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
register(..., geometry=(x_dim, y_dim)) appends a geometry point column
derived from the coordinate dims, synthesized per batch at scan time —
the pivot's coordinate columns already carry the values, so the column
costs an annotation plus (for WKB) a vectorized 21-byte encode of the
rows actually scanned.
Two encodings, chosen per destination:
- "wkb" (default): geoarrow.wkb extension metadata + CRS. DuckDB >=1.2
with spatial loaded ingests the column as GEOMETRY('OGC:CRS84'), so
ST_Within(geometry, ...) works with no ST_Point construction in SQL.
- "point": GeoArrow-native separated coordinates; the struct children
ARE the coordinate arrays (no copy, no parse) for consumers that
execute on native layouts — verified with GeoPandas 1.x
GeoDataFrame.from_arrow, CRS carried through.
Documented sharp edge, measured: engines do not push ST_* functions
into the scan, so a geometry-only predicate defeats chunk pruning and
encodes every row (~29x slower than the paired form on a 10M-row
grid: 101ms bbox vs 2.9s ST_Within-only vs 118ms bbox+ST_Within).
The geospatial docs now state the idiom: bbox conjuncts for pruning,
geometry for exactness.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mmoncadaisla
force-pushed
the
feat/multi-backend-duckdb
branch
from
July 14, 2026 09:37
bcba93a to
f9919b6
Compare
Polars translates float is_in literals imprecisely (reproduced on 1.42: is_in([<stored coordinate>]) matches zero rows). The handle now renders float value lists as OR-chains of degenerate is_between ranges, which compare exactly. Non-float dims keep is_in. A stepped lazy-round-trip window over non-representable float coordinates (linspace(-45, 45, 19)) previously scattered nothing and returned garbage; now row-exact. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The eager path materialized unconditionally: a billion-row result (or a sparse one whose dense coordinate-product grid dwarfs its Arrow payload) exhausted memory rather than erroring. max_result_bytes= now raises a clean ValueError with the running size at both danger points — while collecting the Arrow stream, and before allocating dense arrays (checked against the coordinate product, which is where sparse diagonals blow up). Error before OOM, never truncate; opt-in and unlimited by default. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…pruning The strictness analysis behind count_rows built one guarantee fragment per surviving chunk, capped at 4096 survivors; broader filters fell back to scanning every survivor (a January count on an hourly axis scanned 744 network chunks; step >= 100 over 1M chunks scanned ~1M). Replaced with hierarchical classification: each level buckets the surviving index lists into at most 4096 span-products, decides whole buckets at once with two guarantee-simplification passes (G AND NOT F unsat => proven; G AND F unsat => pruned), and recurses only into mixed cells. Per-chunk coordinate bounds are vectorized (np.minimum.reduceat, cached), so million-chunk axes classify in milliseconds. The prune side also refines the per-dimension pruning with cross-dimension information: paired ranges across dims ((A AND B) OR (C AND D)) no longer read the cross combinations. Measured: count over 1M single-row chunks with a near-universal filter 999,900 rows in 0.21s with ZERO chunk reads (previously scanned every survivor); cross-dim paired ranges read 2 boundary chunks instead of 4; 200 randomized differential checks against numpy ground truth all exact. NaN spans, non-numeric dims and simplifier-rejected expressions still land conservatively in the exact boundary scan. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The previous version used the single-lat-chunk fixture (no cross to refine) and hand-computed the wrong expected count; now differential against numpy on a grid chunked in both dims. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Count-based admission (prefetch=N blocks) under-uses the network with small chunks and overshoots memory with coalesced ones, because block sizes vary. prefetch_bytes= gates admission on estimated pivoted bytes in flight (rows x schema row width), the Lance io_buffer_size semantics; prefetch keeps bounding concurrency (thread count). With a byte budget, raising coalesce_rows no longer multiplies peak RSS. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Engines never push ST_* functions into the scan, so a geometry-only WHERE reads every chunk (documented, measured 29x). This helper renders the bbox range conjuncts from a geometry's envelope (a (xmin, ymin, xmax, ymax) tuple or anything with .bounds — shapely geometries qualify), making the bbox+geometry idiom one f-string instead of hand-copied bounds. Optional pad= margin for ST_DWithin-style predicates. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…mple) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ot streams The two results the re-execution path could not serve now work: spill=True streams the result exactly ONCE with bounded memory into a temporary Parquet file — through the engine handle where one exists (DuckDB spills on its dedicated engine thread, dodging the duckdb-under-worker-threads deadlock entirely; Polars uses its native streaming sink; DataFusion streams batch by batch) or straight from the Arrow stream for one-shot tables/readers — and the ordinary lazy reconstruction then runs over a Polars scan of that file, whose per-window predicates get Parquet row-group pruning. The temp file is deleted when the returned Dataset is garbage collected. Trade-off vs re-execution, by design: one full pass plus temporary disk instead of pay-per-window — the right shape when most of the result will be touched; Polars/DataFusion re-execution remains the default for window-at-a-time access. Verified: the previously deadlocking DuckDB chunked compute topology now passes 6/6 stress runs with repeated full computes; the source is read exactly once (chunk counter), and windows read only the spill file. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DataFusion's section now carries the same relation note DuckDB's has: zarr-datafusion is the engine-native path for plain Zarr, the adapter covers everything else xarray opens plus the round-trip. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A fictional adapter modeled on the DuckDB one: matches by type inspection, register through XarrayPushdownDataset, dispatch via the register_adapter decorator, and the dataset-protocol vs plain-stream trade-off stated. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
duckdb and polars floors (and their version-floor comment) were duplicated between the engine extras and the test extra; the test extra now references xarray_sql[duckdb,polars], so a floor bump has one home. Lockfile regenerated. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ckdb # Conflicts: # README.md # pyproject.toml # tests/test_df.py # xarray_sql/df.py
Four were introduced by the recent benchmark-harness work (untyped result dict, GzipFile into tarfile.open); the rest predate it on the branch (Optional narrowing mypy cannot see through subscripts and closures, Any-typed returns from pyarrow calls). No behavior change; narrowing is expressed with locals and the pyarrow returns are wrapped in their declared types. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A zero batch_size never advances the zero-column scan's row loop (IndexError today, an infinite loop without it); negative values break batch synthesis. Fail with ValueError when the dataset is built instead of mid-scan. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sections follow the dataset's contract (protocol surface, consumer integrations, projection, pruning/counting, scheduling knobs, re-scannability, lifecycle). The four structurally identical count_rows tests become one parametrized contract table. New pinned conditions from adversarial probing: NaN/NaT coordinate chunks are scanned rather than pruned, filter-only columns are read but not returned, an empty projection with a filter reads only the filter column, concurrent and alternating scans stay exact, degenerate tuning values degrade the schedule but never the answer, and non-positive batch_size fails at construction. Every previous regression pin is kept. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
test_era5_integration.py named the instrument, not the subject; it is now test_arrow_dataset_integration.py, pairing with the contract file the way xee pairs ext_test with ext_integration_test. The suite is a dataset x engine matrix: StoreCase entries with expectations computed from declared cadence plus the store's own coordinates, and an engine registry (DuckDB SQL, Polars expressions, DataFusion SQL) that five of the seven tests parametrize over. New coverage: per-engine windowed exactness against a direct xarray read, projection isolation, and concurrent disjoint queries. Fragment consumers (DataFusion) are flagged in the registry since scanner-level coalescing does not apply to them, and the Polars lazy round-trip skips on polars >= 1.43, whose streaming re-execution regressed from 7s to over 10 minutes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
From adversarial review: a run whose cells all failed still exited 0 (only missing _meta was checked); the SKIPPED marker outranked a nonzero exit status, recording crashed cases as skips; and an interrupted source extraction left a corrupt tree every retry reused. The driver now exits nonzero on any error/timeout cell, trusts exit status over the skip marker, and marks extraction complete only after it finishes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
matches is a TypeGuard and register_adapter's bound makes mypy check each adapter against the protocol. runtime_checkable was unused: dispatch goes through matches, never isinstance. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The connection a caller passes is the connection type it gets back; the cast marks the one point where runtime dispatch erases the type. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Both the DataFusion path and the DuckDB adapter split mixed-dimension Datasets with it; living private in sql.py made the generic backends package import the DataFusion-specific module. Tests move with it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
scanner, _scanner_for_blocks, _batches_with_geometry and _batch_generator now read top-down in call order. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Its six pins live on generalized in test_arrow_dataset_integration.py, parametrized per engine and store case. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The second scan is now a 7-day global window instead of a same-shape second day: stale pruning state shows up only when the shape changes, and windows wider than one day get their read-count pin back. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…#240) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
mkdocstrings ignores #: doc-comments and renders :role:`...` and RST hyperlinks as literal text. Doc-comments become attribute docstrings, cross-references become autorefs links (plain code spans for private and external names), and the one RST hyperlink becomes markdown. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The 7-day global scan ran inside the rescannability test, once per engine; the property it pins belongs to the dataset's own scanner. Rescannability keeps two cheap different-shape queries. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mmoncadaisla
marked this pull request as ready for review
August 7, 2026 12:27
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This branch prototypes xarray-sql as the Xarray ↔ engine translator: the library owns exactly two seams — register (a lazy
xarray.Datasetbecomes a table on an engine's own connection) and round-trip (any engine's Arrow result becomes a labeledxr.Datasetagain). SQL dialects, geometry functions, H3, and optimizers stay with each engine and its extension ecosystem. No transpiler, no unified dialect.flowchart TB A["Zarr · NetCDF · GRIB · GeoTIFF · Earth Engine"] --> B["lazy xarray.Dataset"] B --> C["seam 1 — register<br/>xql.register(con, name, ds) / xql.arrow_dataset(ds)"] C --> D["DuckDB<br/>+ spatial, h3, …"] C --> E["Polars"] C --> F["DataFusion"] C --> G["Dask"] D --> H["seam 2 — round-trip<br/>xql.to_dataset(result, template=ds)"] E --> H F --> H H --> I["labeled xr.Dataset — SQL in, array out"]One object serves every engine:
xql.arrow_dataset(ds)is a realpyarrow.dataset.Dataset(the pattern Lance uses forLanceDataset), so DuckDB registers it, Polars scans it viascan_pyarrow_dataset, DataFusion consumes it viaregister_dataset(through the fragments API), and Dask maps overget_fragments(). Ibis works throughibis.duckdb.from_connectionwith zero code.How the pushdown scan works
flowchart TB Q["engine calls scanner(columns, filter)"] --> P["prune chunks<br/>per-dimension shadow fragments;<br/>Arrow guarantee simplification decides satisfiability"] Q --> J["project<br/>only referenced variables are read"] P --> C["coalesce (opt-in)<br/>merge consecutive surviving chunks<br/>into single reads"] C --> L["prefetch pool<br/>bounded by prefetch (threads)<br/>and prefetch_bytes (memory)"] J --> L L --> X["exact filter<br/>pyarrow applies the pushed expression row-exactly"] X --> R["Arrow batches → engine"]Three properties worth calling out:
FileSystemDatasetfragments carry each chunk's coordinate range as apartition_expression(their paths are never opened). Sound for every predicate shape — equality,OR,IN,NOT.count(*)never scans: unfiltered counts are chunk arithmetic, and coordinate-range counts read at most the boundary chunks.The chunked round-trip is engine-generic
xql.to_dataset(result, chunks=...)reconstructs a query result as a chunked, lazyxr.Dataset— each window re-executes the engine's query narrowed to its coordinate range, which flows back into chunk pruning at the source.spill=Trueprovides the alternative one-pass shape: stream once (bounded memory) into a temporary Parquet file that windows re-execute against.flowchart TB R["xql.to_dataset(result, ...)"] --> K{"chunks=?"} K -- "None (default)" --> E["eager: materialize once<br/>(max_result_bytes= guards the stream<br/>and the dense grid)"] K -- "mapping / auto / inherit" --> SP{"spill=?"} SP -- "False (default)" --> RX["re-execution<br/>Polars & DataFusion results"] SP -- "True / directory" --> SPL["one-pass spill → temp Parquet<br/>the chunked path for DuckDB relations<br/>and one-shot Arrow streams"]coords="template"skips coordinate discovery entirely for full-extent results: on ARCO-ERA5 (1.32M hourly chunks) it builds a lazy view over a 1.37-trillion-row table in ~0.3 s with zero source reads; a one-day window then computes in ~2 s reading only the source chunks under it.Performance
Measured on a public 9-billion-pixel cloud-optimized GeoTIFF and a 10M-row synthetic benchmark (
benchmarks/duckdb_pushdown.py):GROUP BY, native res, full table registereddocs/performance.md)AVGscan (vs v1 stream)to_dataset)Peak scan memory is a contract, not an accident: bounded by
prefetch × pivoted-block-sizeregardless of data scanned (a 772M-row month-scale ARCO-ERA5 aggregation peaks at the same ~0.75 GB RSS as the week-scale scan). The pivot fast path (whole-partition repeat/tile coordinate columns, zero-copy batch slices) and the round-trip fast paths (affine-axis scatter, grid reshape) live in shared code, so the existing DataFusion engine benefits equally.Beyond microbenchmarks,
benchmarks/geospatial/runs nine staples of geospatial/climate analysis (NDVI, climatology, anomalyJOIN, forecast skill vs WeatherBench 2, raster×vector zonal stats, PROJ-UDF reprojection, weight-table regridding, warp) in SQL against real cloud datasets, each asserted against an xarray reference;engine_suite.pyrepeats the portable cases across DataFusion (both the native table provider and the pure-Python pyarrow path, as separatedatafusion/datafusion-arrowengines), DuckDB, and Polars on GCE VMs, provisioning the compiled native module on each VM. The write-up isdocs/geospatial.md.What's in the branch
xarray_sql/backends/— adapter protocol with dispatch on connection type; DataFusion adapter delegates to the existing table provider; DuckDB adapter registers the pushdown dataset (mixed-dimension datasets split into one table per dim group, sharing coordinate reads).xarray_sql/backends/pyarrow.py—XarrayPushdownDataset(projection pushdown, shadow pruning, coalescing, prefetch, fragments API) andXarrayArrowStream(re-scannable C-stream fallback).xarray_sql/roundtrip.py+xarray_sql/lazyscan.py— engine-agnosticxql.to_datasetaccepting DuckDB relations, Polars frames, DataFusion DataFrames,pyarrowtables/readers, or any__arrow_c_stream__object; eager, re-executing, and spill-backed chunked reconstruction; metadata recovery from a template Dataset.xarray_sql/geometry.py—register(..., geometry=("x", "y"))derives a GeoArrow pointgeometrycolumn (WKB for DuckDB-nativeGEOMETRY, or GeoArrow-separated for GeoPandas/lonboard), CRS tagged;xql.bbox_conjunctsrenders prunable bbox predicates from any geometry's envelope.df.py/ds.py.docs/engines.md(the engine model, per-engine usage, support matrix),docs/performance.md(measured tuning guide),docs/limitations.md(known issues, pinned by tests),docs/geospatial.md(the relational-operations write-up with the benchmark results), reworkeddocs/examples.mdand README.[duckdb](duckdb>=1.4) and[polars](polars>=1.33) extras.Testing
275 tests, all green. Beyond unit coverage:
OR/IN, filter-column-outside-projection,LIMIT, empty results, NULL semantics) passes identically on duckdb 1.4.5 LTS and 1.5.4, plus the Polars battery on 1.42.max_result_bytesenforcement on eager collection and dense allocation.ChunkedArrayslipping into the pivot's fast path — all fixed with regression tests.Known limitations / follow-ups
chunks=raises immediately): re-execution from worker threads deadlocks inside duckdb-python 1.4–1.5 on CPython 3.12, sospill=Trueis the chunked path there. Documented with the full story indocs/limitations.md.XarrayPushdownDatasetsubclasses pyarrow's cythonDatasetwithout initializing the native base (as Lance does); dangerous inherited members are stubbed, and the contract is pinned by tests — re-verify on pyarrow upgrades.is_infloat literals imprecisely (reproducible without xarray-sql); documented, use range predicates. Polars also has no geometry types — thegeometrycolumn arrives as plain binary there.🤖 Generated with Claude Code